package mill.javalib

import mainargs.Flag
import mill.*
import mill.api.{PathRef, Task, *}
import mill.api.daemon.internal.internal
import mill.javalib.CoursierModule.Resolver
import mill.javalib.api.JvmWorkerUtil.{isBinaryBridgeAvailable, isDotty, isDottyOrScala3}
import mill.javalib.api.internal.InternalJvmWorkerApi as InternalJvmWorkerApi
import mill.javalib.api.{JvmWorkerApi, JvmWorkerUtil, Versions}
import mill.javalib.api.{JvmWorkerArgs}
import mill.javalib.api.internal.{ZincCompilerBridgeProvider}

/**
 * A default implementation of [[JvmWorkerModule]]
 */
object JvmWorkerModule extends ExternalModule with JvmWorkerModule with CoursierModule {
  lazy val millDiscover = Discover[this.type]
}

/**
 * A module managing an in-memory Zinc Scala incremental compiler
 */
trait JvmWorkerModule extends OfflineSupportModule with CoursierModule {

  def classpath: T[Seq[PathRef]] = Task {
    defaultResolver().classpath(Seq(
      Dep.millProjectModule("mill-libs-javalib-worker")
    ))
  }

  /**
   * This is actually the testrunner classpath, not the scalalib classpath,
   * but the wrong name is preserved for backwards compatibility
   */
  def scalalibClasspath: T[Seq[PathRef]] = Task {
    defaultResolver().classpath(Seq(
      Dep.millProjectModule("mill-libs-javalib-testrunner")
    ))
  }

  def testrunnerEntrypointClasspath: T[Seq[PathRef]] = Task {
    defaultResolver().classpath(Seq(
      Dep.millProjectModule("mill-libs-javalib-testrunner-entrypoint", artifactSuffix = "")
    ))
  }

  def backgroundWrapperClasspath: T[Seq[PathRef]] = Task {
    defaultResolver().classpath(Seq(
      Dep.millProjectModule("mill-libs-javalib-backgroundwrapper", artifactSuffix = "")
    ))
  }

  /** Whether Zinc debug logging is enabled. */
  def zincLogDebug: T[Boolean] = Task.Input(Task.ctx().log.debugEnabled)

  def worker: Worker[JvmWorkerApi] = internalWorker
  def internalWorkerClassLoader = Task.Worker {
    mill.util.Jvm.createClassLoader(classpath().map(_.path), getClass.getClassLoader)
  }

  @internal def internalWorker: Worker[InternalJvmWorkerApi] = Task.Worker {
    val jobs = Task.ctx().jobs

    val cl = internalWorkerClassLoader()

    val ctx = Task.ctx()
    val zincCompilerBridge = ZincCompilerBridgeProvider(
      workspace = ctx.dest,
      logInfo = ctx.log.info,
      acquire = (scalaVersion, scalaOrganization) =>
        scalaCompilerBridgeJarV2(scalaVersion, scalaOrganization, defaultResolver()).map(_.path)
    )

    val args = JvmWorkerArgs(
      zincCompilerBridge,
      classPath = classpath().map(_.path),
      jobs = jobs,
      zincLogDebug = zincLogDebug(),
      close0 = () => cl.close()
    )

    cl.loadClass("mill.javalib.worker.JvmWorkerImpl")
      .getConstructor(classOf[JvmWorkerArgs])
      .newInstance(args)
      .asInstanceOf[InternalJvmWorkerApi]
  }

  private[mill] def scalaCompilerBridgeJarV2(
      scalaVersion: String,
      scalaOrganization: String,
      resolver: Resolver
  )(using ctx: TaskCtx): ZincCompilerBridgeProvider.AcquireResult[PathRef] = {
    val (scalaVersion0, scalaBinaryVersion0) = scalaVersion match {
      case _ => (scalaVersion, JvmWorkerUtil.scalaBinaryVersion(scalaVersion))
    }

    val (bridgeDep, bridgeName, bridgeVersion) =
      if (isDottyOrScala3(scalaVersion0)) {
        val org = scalaOrganization
        val name =
          if (isDotty(scalaVersion0)) "dotty-sbt-bridge"
          else "scala3-sbt-bridge"
        val version = scalaVersion
        (mvn"$org:$name:$version", name, version)
      } else if (JvmWorkerUtil.millCompilerBridgeScalaVersions.contains(scalaVersion0)) {
        val org = "com.lihaoyi"
        val name = s"mill-scala-compiler-bridge_$scalaVersion"
        val version = Versions.millCompilerBridgeVersion
        (mvn"$org:$name:$version", name, version)
      } else {
        val org = "org.scala-sbt"
        val name = "compiler-bridge"
        val version = Versions.zinc
        (
          mvn"$org:${name}_${scalaBinaryVersion0}:$version",
          s"${name}_$scalaBinaryVersion0",
          version
        )
      }

    val useSources = !isBinaryBridgeAvailable(scalaVersion)

    val deps = resolver.classpath(
      Seq(bridgeDep.bindDep("", "", "")),
      sources = useSources,
      mapDependencies = Some(overrideScalaLibrary(scalaVersion, scalaOrganization))
    )

    val bridgeJar = JvmWorkerUtil.grepJar(deps, bridgeName, bridgeVersion, useSources)

    if (useSources) {
      val classpath = compilerInterfaceClasspath(scalaVersion, scalaOrganization, resolver)
      ZincCompilerBridgeProvider.AcquireResult.NotCompiled(classpath, bridgeJar)
    } else ZincCompilerBridgeProvider.AcquireResult.Compiled(bridgeJar)
  }

  @deprecated("This is an internal API that has been accidentally exposed.", "Mill 1.0.2")
  def scalaCompilerBridgeJar(
      scalaVersion: String,
      scalaOrganization: String,
      resolver: Resolver
  )(using ctx: TaskCtx): (Option[Seq[PathRef]], PathRef) =
    scalaCompilerBridgeJarV2(
      scalaVersion = scalaVersion,
      scalaOrganization = scalaOrganization,
      resolver
    ) match {
      case ZincCompilerBridgeProvider.AcquireResult.Compiled(bridgeJar) => (None, bridgeJar)
      case ZincCompilerBridgeProvider.AcquireResult.NotCompiled(classpath, bridgeSourcesJar) =>
        (Some(classpath), bridgeSourcesJar)
    }

  def compilerInterfaceClasspath(
      scalaVersion: String,
      scalaOrganization: String,
      resolver: Resolver
  )(using ctx: TaskCtx): Seq[PathRef] = {
    resolver.classpath(
      deps = Seq(mvn"org.scala-sbt:compiler-interface:${Versions.zinc}".bindDep("", "", "")),
      // Since Zinc 1.4.0, the compiler-interface depends on the Scala library
      // We need to override it with the scalaVersion and scalaOrganization of the module
      mapDependencies = Some(overrideScalaLibrary(scalaVersion, scalaOrganization))
    )
  }

  def overrideScalaLibrary(
      scalaVersion: String,
      scalaOrganization: String
  )(dep: coursier.Dependency): coursier.Dependency = {
    if (dep.module.name.value == "scala-library") {
      dep.withModule(dep.module.withOrganization(coursier.Organization(scalaOrganization)))
        .withVersion(scalaVersion)
    } else dep
  }

  override def prepareOffline(all: Flag): Command[Seq[PathRef]] = Task.Command {
    (super.prepareOffline(all)() ++ classpath()).distinct
  }

  // noinspection ScalaUnusedSymbol - Task.Command
  def prepareOfflineCompiler(scalaVersion: String, scalaOrganization: String): Command[Unit] =
    Task.Command {
      classpath()
      scalaCompilerBridgeJarV2(scalaVersion, scalaOrganization, defaultResolver())
      ()
    }

}
